{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing required libraries\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "# Create a directed graph\n",
    "G = nx.DiGraph()\n",
    "\n",
    "# Adding nodes\n",
    "G.add_nodes_from([\"Start\", \"Decision\", \"Action 1\", \"Action 2\", \"End\"])\n",
    "\n",
    "# Adding edges\n",
    "G.add_edges_from([(\"Start\", \"Decision\"),\n",
    "                  (\"Decision\", \"Action 1\"),\n",
    "                  (\"Decision\", \"Action 2\"),\n",
    "                  (\"Action 1\", \"End\"),\n",
    "                  (\"Action 2\", \"End\")])\n",
    "\n",
    "labels = {\n",
    "    (\"Start\", \"Decision\"): \"Begin\",\n",
    "    (\"Decision\", \"Action 1\"): \"Yes\",\n",
    "    (\"Decision\", \"Action 2\"): \"No\",\n",
    "    (\"Action 1\", \"End\"): \"\",\n",
    "    (\"Action 2\", \"End\"): \"\",\n",
    "}\n",
    "\n",
    "node_colors = ['skyblue' if n == 'Decision' else 'lightgray' for n in G.nodes]\n",
    "\n",
    "# nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=2000, font_size=10, font_color='black', font_weight='bold', arrows=True)\n",
    "# nx.draw_networkx_edge_labels(G, pos, edge_labels=labels)\n",
    "\n",
    "# Correcting the error and drawing the graph again with rounded corners for the rectangles\n",
    "\n",
    "# Create a new figure and axis for larger size\n",
    "\n",
    "pos = {\n",
    "    \"Start\": (0, 2),\n",
    "    \"Decision\": (1, 2),\n",
    "    \"Action 1\": (1.5, 3),\n",
    "    \"Action 2\": (1.5, 1),\n",
    "    \"End\": (2, 2)\n",
    "}\n",
    "\n",
    "# nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=5000, font_size=10, font_color='black', font_weight='bold', node_shape='s', arrows=True)\n",
    "# nx.draw_networkx_edge_labels(G, pos, edge_labels=labels)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-23 23:26:21.703 | INFO     | sweepai.config.server:<module>:195 - Using Sandbox URL: http://24.199.78.105:8081\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "378b4e13-f0fd-46a4-9818-49bc9f23d55f\n",
      "409a3705-50bc-405c-9aef-d49daaf2aa62\n"
     ]
    }
   ],
   "source": [
    "from sweepai.core.entities import FileChangeRequest\n",
    "\n",
    "first_unit_test = FileChangeRequest(\n",
    "    filename=\"a_test.py\",\n",
    "    instructions=\"Write unit tests\",\n",
    "    change_type=\"create\",\n",
    "    status=\"succeeded\"\n",
    ")\n",
    "\n",
    "first_check = FileChangeRequest(\n",
    "    filename=\"a_test.py\", \n",
    "    instructions=\"Write unit tests\",\n",
    "    change_type=\"check\",\n",
    "    parent=first_unit_test,\n",
    "    status=\"failed\"\n",
    ")\n",
    "\n",
    "first_fix = FileChangeRequest(\n",
    "    filename=\"a.py\",\n",
    "    instructions=\"Fix failed unit test\",\n",
    "    change_type=\"modify\",\n",
    "    parent=first_check,\n",
    "    status=\"running\"\n",
    ")\n",
    "\n",
    "file_change_requests = [\n",
    "    first_unit_test,\n",
    "    first_check,\n",
    "    first_fix,\n",
    "    FileChangeRequest(\n",
    "        filename=\"a_test.py\", \n",
    "        instructions=\"Fix failed unit test\",\n",
    "        change_type=\"check\",\n",
    "        parent=first_fix\n",
    "    ),\n",
    "    FileChangeRequest(\n",
    "        filename=\"README.md\", \n",
    "        instructions=\"Write unit tests\",\n",
    "        change_type=\"modify\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "# print(file_change_requests)\n",
    "print(first_check.id_)\n",
    "print(file_change_requests[3].parent.id_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.43.0 (0)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"466pt\" height=\"408pt\"\n viewBox=\"0.00 0.00 466.04 408.48\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(36 372.48)\">\n<title>%3</title>\n<polygon fill=\"dimgray\" stroke=\"transparent\" points=\"-36,36 -36,-372.48 430.04,-372.48 430.04,36 -36,36\"/>\n<!-- start -->\n<g id=\"node1\" class=\"node\">\n<title>start</title>\n</g>\n<!-- fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043 -->\n<g id=\"node2\" class=\"node\">\n<title>fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043</title>\n<ellipse fill=\"#0ee832\" stroke=\"black\" cx=\"109.5\" cy=\"-309.61\" rx=\"56.64\" ry=\"26.74\"/>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-313.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">✓ Create</text>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-298.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">a_test.py</text>\n</g>\n<!-- start&#45;&gt;fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043 -->\n<g id=\"edge5\" class=\"edge\">\n<title>start&#45;&gt;fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1.2,-309.61C4.87,-309.61 22.63,-309.61 42.74,-309.61\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"42.75,-313.11 52.75,-309.61 42.75,-306.11 42.75,-313.11\"/>\n<text text-anchor=\"middle\" x=\"26.97\" y=\"-316.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Start</text>\n</g>\n<!-- acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4 -->\n<g id=\"node3\" class=\"node\">\n<title>acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"320.5\" cy=\"-309.61\" rx=\"73.58\" ry=\"26.74\"/>\n<text text-anchor=\"middle\" x=\"320.5\" y=\"-313.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">▶ Modify</text>\n<text text-anchor=\"middle\" x=\"320.5\" y=\"-298.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">README.md</text>\n</g>\n<!-- fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043&#45;&gt;acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4 -->\n<g id=\"edge4\" class=\"edge\">\n<title>fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043&#45;&gt;acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4</title>\n<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M166.32,-309.61C188,-309.61 213.2,-309.61 236.86,-309.61\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"236.89,-313.11 246.89,-309.61 236.89,-306.11 236.89,-313.11\"/>\n<text text-anchor=\"middle\" x=\"206.51\" y=\"-316.41\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Continue</text>\n</g>\n<!-- 378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f -->\n<g id=\"node4\" class=\"node\">\n<title>378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f</title>\n<polygon fill=\"#ee0000\" stroke=\"black\" points=\"149.5,-231.74 69.5,-231.74 69.5,-193.74 149.5,-193.74 149.5,-231.74\"/>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-216.54\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">✗ Check</text>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-201.54\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">a_test.py</text>\n</g>\n<!-- fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043&#45;&gt;378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f -->\n<g id=\"edge1\" class=\"edge\">\n<title>fff516f5&#45;94d4&#45;412f&#45;9edd&#45;0cf13231a043&#45;&gt;378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f</title>\n<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M109.5,-282.69C109.5,-270.09 109.5,-254.91 109.5,-242.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"113,-241.82 109.5,-231.82 106,-241.82 113,-241.82\"/>\n<text text-anchor=\"middle\" x=\"162.5\" y=\"-253.54\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Check changes</text>\n</g>\n<!-- end -->\n<g id=\"node6\" class=\"node\">\n<title>end</title>\n</g>\n<!-- acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4&#45;&gt;end -->\n<g id=\"edge6\" class=\"edge\">\n<title>acd31770&#45;10b8&#45;4d49&#45;b857&#45;43b4d62a28d4&#45;&gt;end</title>\n<path fill=\"none\" stroke=\"black\" d=\"M320.5,-282.49C320.5,-227.98 320.5,-102.61 320.5,-47.19\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"324,-47.08 320.5,-37.08 317,-47.08 324,-47.08\"/>\n<text text-anchor=\"middle\" x=\"334.5\" y=\"-164.54\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">End</text>\n</g>\n<!-- 409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62 -->\n<g id=\"node5\" class=\"node\">\n<title>409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62</title>\n<ellipse fill=\"yellow\" stroke=\"black\" cx=\"109.5\" cy=\"-115.87\" rx=\"57.05\" ry=\"26.74\"/>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-119.67\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">⋯ Modify</text>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-104.67\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">a.py</text>\n</g>\n<!-- 378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f&#45;&gt;409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62 -->\n<g id=\"edge2\" class=\"edge\">\n<title>378b4e13&#45;f0fd&#45;46a4&#45;9818&#45;49bc9f23d55f&#45;&gt;409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62</title>\n<path fill=\"none\" stroke=\"black\" d=\"M109.5,-193.45C109.5,-182.03 109.5,-166.95 109.5,-153.18\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"113,-153 109.5,-143 106,-153 113,-153\"/>\n<text text-anchor=\"middle\" x=\"189.5\" y=\"-164.54\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">More changes required</text>\n</g>\n<!-- efa253dd&#45;93f6&#45;4859&#45;920c&#45;46e8812b6dc0 -->\n<g id=\"node7\" class=\"node\">\n<title>efa253dd&#45;93f6&#45;4859&#45;920c&#45;46e8812b6dc0</title>\n<polygon fill=\"white\" stroke=\"black\" points=\"149.5,-38 69.5,-38 69.5,0 149.5,0 149.5,-38\"/>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-22.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">▶ Check</text>\n<text text-anchor=\"middle\" x=\"109.5\" y=\"-7.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">a_test.py</text>\n</g>\n<!-- 409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62&#45;&gt;efa253dd&#45;93f6&#45;4859&#45;920c&#45;46e8812b6dc0 -->\n<g id=\"edge3\" class=\"edge\">\n<title>409a3705&#45;50bc&#45;405c&#45;9aef&#45;d49daaf2aa62&#45;&gt;efa253dd&#45;93f6&#45;4859&#45;920c&#45;46e8812b6dc0</title>\n<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M109.5,-88.95C109.5,-76.35 109.5,-61.17 109.5,-48.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"113,-48.08 109.5,-38.08 106,-48.08 113,-48.08\"/>\n<text text-anchor=\"middle\" x=\"162.5\" y=\"-59.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Check changes</text>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7fc2a013ee90>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from graphviz import Digraph\n",
    "\n",
    "dot = Digraph(comment='FileChangeRequest Tree')\n",
    "fontname = \"helvetica\"\n",
    "fontsize = \"12\"\n",
    "dot.attr(fontname=fontname)\n",
    "dot.attr(title=\"Tree\\n\\n\")\n",
    "dot.attr(pad=\"0.5\")\n",
    "dot.attr(bgcolor='dimgray')\n",
    "\n",
    "dot.edge_attr.update(color='lightgray')\n",
    "fontcolor = \"black\"\n",
    "\n",
    "\n",
    "ranks = {}\n",
    "\n",
    "for i, fcr in enumerate(file_change_requests):\n",
    "    if fcr.parent is None:\n",
    "        ranks[fcr.id_] = 0\n",
    "    else:\n",
    "        ranks[fcr.id_] = ranks[fcr.parent.id_] + 1\n",
    "\n",
    "for layer in range(max(ranks.values()) + 1):\n",
    "    with dot.subgraph() as c:\n",
    "        if layer == 0:\n",
    "            c.attr(label=\"Original plan\", labelloc=\"t\", labeljust=\"l\", rank=\"same\", fontname=fontname, fontcolor=fontcolor)\n",
    "            c.node('start', '', shape='none', width=\"0\", fontname=fontname, fontcolor=fontcolor)\n",
    "        else:\n",
    "            c.attr(label=f\"Layer {layer}\", rank=\"same\", fontname=fontname, fontcolor=fontcolor)\n",
    "        if layer == max(ranks.values()):\n",
    "            c.node('end', '', shape='none', width=\"0\", fontname=fontname, fontcolor=fontcolor)\n",
    "        for fcr in file_change_requests:\n",
    "            if ranks[fcr.id_] == layer:\n",
    "                if fcr.change_type == \"check\":\n",
    "                    c.node(fcr.id_, fcr.summary, shape=\"rectangle\", fillcolor=fcr.color, style=\"filled\", fontname=fontname, fontcolor=fontcolor, color=fontcolor)\n",
    "                else:\n",
    "                    c.node(fcr.id_, fcr.summary, fillcolor=fcr.color, style=\"filled\", fontname=fontname, fontcolor=fontcolor, color=fontcolor)\n",
    "\n",
    "last_item_per_layer = {layer: None for layer in range(max(ranks.values()) + 1)}\n",
    "\n",
    "for fcr in file_change_requests:\n",
    "    if fcr.parent:\n",
    "        if fcr.change_type == \"check\":\n",
    "            dot.edge(fcr.parent.id_, fcr.id_, label=\"Check changes\", style=\"dashed\", fontname=fontname, fontcolor=fontcolor, color=fontcolor)\n",
    "        elif fcr.parent.change_type == \"check\":\n",
    "            dot.edge(fcr.parent.id_, fcr.id_, label=\"More changes required\", fontname=fontname, fontcolor=fontcolor, color=fontcolor)\n",
    "        else:\n",
    "            dot.edge(fcr.parent.id_, fcr.id_)\n",
    "    elif last_item_per_layer[ranks[fcr.id_]] is not None:\n",
    "        dot.edge(last_item_per_layer[ranks[fcr.id_]].id_, fcr.id_, label=\"Continue\", style=\"dashed\", fontname=fontname, fontcolor=fontcolor, color=fontcolor)\n",
    "    last_item_per_layer[ranks[fcr.id_]] = fcr\n",
    "\n",
    "dot.edge('start', file_change_requests[0].id_, label=\"Start\", fontname=fontname, color=fontcolor)\n",
    "dot.edge(file_change_requests[-1].id_, 'end', label=\"End\", fontname=fontname, color=fontcolor)\n",
    "\n",
    "dot.format = \"svg\"\n",
    "dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 ('sweepai-u_CIt3kb-py3.10')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "25d341f3248a096a89b9dbf6eec8e41f63aed02f6ba059df22a49224e3e8f1b0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
